package net.csdn.business.gateway.filter;

import com.alibaba.fastjson2.JSONObject;
import lombok.extern.slf4j.Slf4j;
import net.csdn.business.common.constant.GatewayConstants;
import net.csdn.business.common.utils.IpUtil;
import net.csdn.business.gateway.component.CustomRedisRateLimiter;
import net.csdn.business.gateway.component.HourRateLimitRedisLua;
import net.csdn.business.gateway.config.RateLimitConfig;
import net.csdn.business.gateway.utils.ResponseUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;

import java.time.Instant;
import java.time.temporal.ChronoUnit;

/**
 * @Auther: zhangyalei
 * @Date: 2022/12/22 10:52
 * @Description: 自定义限流过滤器
 */
@Slf4j
@Component
public class CustomRateLimitGatewayFilterFactory extends AbstractGatewayFilterFactory {

    @Autowired
    private CustomRedisRateLimiter myRedisRateLimiter;

    @Autowired
    private HourRateLimitRedisLua hourRateLimitRedisLua;

    @Autowired
    private RateLimitConfig rateLimitConfig;


    @Override
    public GatewayFilter apply(Object config) {
        return (exchange, chain) -> {
            boolean flageS=false; //每秒速率限制是否通过标志
            boolean flageH=false; //每秒小时速率限制是否通过标志
            HttpHeaders headers = exchange.getRequest().getHeaders();
            String identityId=headers.getFirst(GatewayConstants.IDENTITY_ID);
            String identityType=headers.getFirst(GatewayConstants.IDENTITY_TYPE);
            String token=headers.getFirst(GatewayConstants.TOKEN);
            if(StringUtils.isNotBlank(token)&&StringUtils.isBlank(identityId)){
                //特殊的jwt-token(用在应用的某些特殊的openApi上),不进行限流检测
                return chain.filter(exchange);
            }
            String desc=null;
            JSONObject hourRes=null;
            Integer total=null;
            Integer remain=null;//剩余次数
            Long time=null;//剩余重置时间
            String ip = IpUtil.getIp(exchange.getRequest());
            //对于普通用户的限流策略
            if(StringUtils.isNotBlank(identityId)){
                String hourKey=GatewayConstants.RATE_LIMIE_NAME_SPACE_H+identityId;
                switch (identityType){
                    case "USER":
                        if(userRateLimiter(identityId)){
                            flageS=true;
                        }
                        //限流通过，则校验每小时的
                        total=rateLimitConfig.getUidHour();
                        hourRes=hourRateLimitRedisLua.subtracts(hourKey,3600,total);
                        break;
                    case "APP":
                        if(appRateLimiter(identityId)){
                            flageS=true;
                        }
                        //限流通过，则校验每小时的
                        total=rateLimitConfig.getAppHour();
                        hourRes=hourRateLimitRedisLua.subtracts(hourKey,3600,total);
                        break;
                }
                desc="API rate limit exceeded,documentation_url: https://docs.atomgit.com";
            }else{
                if(StringUtils.isNotBlank(ip)){
                    if(ipRateLimiter(ip)){
                        flageS=true;
                    }
                    String key=GatewayConstants.RATE_LIMIE_NAME_SPACE_H+ip;
                    total=rateLimitConfig.getIpHour();
                    hourRes=hourRateLimitRedisLua.subtracts(key,3600,total);
                    desc=String.format("API rate limit exceeded for %s. (But here's the good news: Authenticated requests get a higher rate limit. Check out the documentation for more details. documentation_url: 'https://docs.atomgit.com'",ip);
                }
            }
            if(hourRes!=null){
                remain=hourRes.getInteger("remain");//剩余次数
                Integer seconds=hourRes.getInteger("time");//剩余次数
                time= Instant.now().plus(seconds, ChronoUnit.SECONDS).getEpochSecond();
                Integer hourflage=hourRes.getInteger("flag");
                flageH=hourflage==1?true:false;
            }
            //校验每小时呢
            Integer finalUsed=total-remain;
            //获取响应内容
      /*      //修改返回Response
            Integer finalTotal = total;
            Integer finalRemain = remain;
            Integer finalUsed=total-remain;
            Integer finalTime = time;
            ServerHttpResponseDecorator responseDecorator = new ServerHttpResponseDecorator(exchange.getResponse()) {

                @Override
                public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                    HttpHeaders httpHeaders = new HttpHeaders();
                    httpHeaders.add(GatewayConstants.X_RATE_LIMIT, finalTotal.toString());
                    httpHeaders.add(GatewayConstants.X_RATE_LIMIT_USED, finalUsed.toString());
                    httpHeaders.add(GatewayConstants.X_RATE_LIMIT_REMAINING, finalRemain.toString());
                    httpHeaders.add(GatewayConstants.X_RATE_LIMIT_RESET, finalTime.toString());
                    Integer statusCode=exchange.getResponse().getStatusCode().value();
                    //直接放行
                    if(statusCode!=200){
                        return chain.filter(exchange);
                    }

                    ClientResponse clientResponse = ClientResponse
                            .create(exchange.getResponse().getStatusCode())
                            .headers(headers -> headers.putAll(httpHeaders))
                            .body(Flux.from(body)).build();

                    //修改body
                    Mono<String> modifiedBody = clientResponse.bodyToMono(String.class);


                    BodyInserter bodyInserter = BodyInserters.fromPublisher(modifiedBody,
                            String.class);

                    CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(
                            exchange, exchange.getResponse().getHeaders());
                    return bodyInserter.insert(outputMessage, new BodyInserterContext())
                            .then(Mono.defer(() -> {
                                Flux<DataBuffer> messageBody = outputMessage.getBody();
                                HttpHeaders headers = getDelegate().getHeaders();
                                if (!headers.containsKey(HttpHeaders.TRANSFER_ENCODING)) {
                                    messageBody = messageBody.doOnNext(data -> headers
                                            .setContentLength(data.readableByteCount()));
                                }
                                return getDelegate().writeWith(messageBody);
                            }));
                }

                @Override
                public Mono<Void> writeAndFlushWith(
                        Publisher<? extends Publisher<? extends DataBuffer>> body) {
                    return writeWith(Flux.from(body).flatMapSequential(p -> p));
                }

            };*/
            exchange.getResponse().getHeaders().add(GatewayConstants.X_RATE_LIMIT, total.toString());
            exchange.getResponse().getHeaders().add(GatewayConstants.X_RATE_LIMIT_USED, finalUsed.toString());
            exchange.getResponse().getHeaders().add(GatewayConstants.X_RATE_LIMIT_REMAINING, remain.toString());
            exchange.getResponse().getHeaders().add(GatewayConstants.X_RATE_LIMIT_RESET, time.toString());
            //exchange.getResponse().getHeaders().add(CommonConstants.FIELD_EAGLEEYE_TRACEID, MDC.get(CommonConstants.TRACE_ID));
            //ServerWebExchange webExchange = exchange.mutate().response(responseDecorator).build();
            if(!flageS||!flageH){
                return ResponseUtils.webFluxResponseWriter(exchange.getResponse(), HttpStatus.FORBIDDEN,"forbidden",desc);
            }
            return chain.filter(exchange);
        };
    }








    //按照ip，每秒限流
    private Boolean ipRateLimiter(String ip){
        //从数据库中获取到该IP地址对应的限流参数
        return myRedisRateLimiter.isAllowed("ip:"+ip,
                rateLimitConfig.getIpSecond(), //令牌桶每秒填充平均速率
                rateLimitConfig.getIpSecond()); //令牌桶总容量
    }
    //按照uid，每秒限流
    private Boolean userRateLimiter(String userId){
        //获取到调用客户端的IP地址
        return myRedisRateLimiter.isAllowed("uid:"+userId+":",
                rateLimitConfig.getUidSecond(), //令牌桶每秒填充平均速率
                rateLimitConfig.getUidSecond()); //令牌桶总容量
    }

    //按照APP的id，每秒限流
    private Boolean appRateLimiter(String appId){
        //获取到调用客户端的IP地址
        return myRedisRateLimiter.isAllowed("app:"+appId+":",
                rateLimitConfig.getAppSecond(), //令牌桶每秒填充平均速率
                rateLimitConfig.getAppSecond()); //令牌桶总容量
    }

}
